#include "trm/segment.h"
#include "trm/reference_frame.h"
#include "trm/span_arithm.h"
#include "trm/mincut.h"
//#include "refine/smooth2.h"
//#include "refine/smooth_b2.h"
#include "refine/smooth_c.h"
//#include "refine/smooth_d.h"
//#include "refine/smooth_e.h"
//#include "refine/smooth_ea.h"
#include "refine/smooth_despill_a.h"
#include "refine/smooth_despill_b.h"
//#include "refine/smooth_despill_c.h"
#include "refine/bit_mask.h"
#include "refine/refine2.h"
//#include "refine/refine3.h"
#include "outline/decrease_mask_freedom.h"
#include "xtd/redirect.h"
#include "imglib/bmpmorph.h"

namespace trm {

std::string debug_dir()
{
	const char * env = getenv("PM_DEBUG_DIR");
	if (env == 0 || *env == 0)
		return std::string();
	std::string path = env;
	if (*path.rbegin() != '\\' && *path.rbegin() != '/')
		path.push_back('\\');
	return path;
}

class debug_log : public std_log_redirect
{
public:
	debug_log() {
		std::string path = debug_dir();
		if (path.empty()) 
			redirect((std::streambuf*)0);
		else
			redirect((path+"trm.log").c_str());
	}
};

static void wool_removal(trm::bit_mask & m, int d)
{
	//new_m = (src_mask & open(m)) | (~src_mask & close(m))
	trm::bit_mask o, c;
	img::open4(o, m, d);
	img::combine(o, m, img::bit_and());
	img::close4(c, m, d);
	img::combine(c, m, img::bit_and_not());
	img::combine(o, c, img::bit_or());
	m.swap(o);
}

#define TestCancel() { if (is_canceled()) return; }

struct segment::calc_data
{
	calc_data(segment & s, span const & r) : request(r), ref(s) { }
	span const & request;
	direction dir;
	reference_frame ref; ///<it may happen that origin frame will be reference for both backward and forward directions
};

void segment::calc(bool key_mask_refined, span const & request)
{
	assert(_accessible_span.inside(_origin));
	assert(_accessible_span.inside(request));
	TestCancel();
	debug_log dlog;
	mswin::timing_mark mark(_timing, "calculating");

	bool need_origin = !has_origin_masks() || origin_masks()->b.empty();
	bool need_adjust = need_origin && !key_mask_refined && 
		(_param.mincut || _param.refine_mask || _param.refine2_mask /*|| _param.refine3_mask*/ || _param.decrease_mask_freedom);

	//this is temporary workaround for badly saved projects
	if (!need_origin && _param.decrease_mask_freedom && origin_masks()->contour.empty())
		need_origin = need_adjust = true;

	if (need_origin) {
		trim(span(_origin, _origin));
		twomasks_sptr p(new twomasks);
		p->b = _key_mask;
		if (need_adjust) {
			icache::cached_image_sptr iptr = load_origin_frame(request);
			if (_param.mincut) {
				TestCancel();
				mswin::timing_mark mark(_timing, "mincut segmenting");
				mincut::params mp;
				mp.pyramid_levels = 0;
				mp.max_motion = 2; //adjust user-drawn mask maximum on a couple of pixels //_param.mincut_max_motion, 
				mp.hw = _param.mincut_max_link_distance;
				mp.nn = _param.mincut_link_count;
				mp.aspect = _param.mincut_aspect;
				mp.source = mincut::source_target;
				mp.gray_levels_exp = _param.mincut_antialiasing_levels_exp;
				mp.timing_tree = &_timing;
				trm::grey_mask g;
				mincut::adjust(iptr->im, p->b, 
					&p->b, 0, mp,
					_param.mincut_and_antialiasing() ? &g : 0);
				if (!g.empty())
					p->g.store(g);
			}
			if (_param.refine_mask) {
				TestCancel();
				mswin::timing_mark mark(_timing, "refining");
				refine::bit_mask::params rp;
				rp.whs = _param.refine_window_half_size;
				rp.auto_t = _param.refine_auto;
				rp.t = _param.refine_enough_statistics;
				rp.prefer = _param.refine_preference;
				refine::bit_mask::adjust(iptr->im, p->b, rp);
			}
			if (_param.refine2_mask) {
				TestCancel();
				mswin::timing_mark mark(_timing, "refining2");
				using namespace refine::refine2;
				params rp;
				rp.max_motion = _param.refine2_max_motion;
				rp.max_level = _param.refine2_max_level;
				track(iptr->im, p->b, iptr->im, p->b, rp);
			}
			/*if (_param.refine3_mask) {
				TestCancel();
				mswin::timing_mark mark(_timing, "refining3");
				using namespace refine::refine3;
				params rp;
				rp.t = _param.refine3_enough_statistics;
				rp.prefer = _param.refine3_preference;
				adjust(iptr->im, p->b, rp);
			}*/
			if (_param.wool_removal > 0) {
				TestCancel();
				wool_removal(p->b, _param.wool_removal);
			}
			if (_param.decrease_mask_freedom) {
				TestCancel();
				decrease_mask_freedom(*p);
			}
		}
		set_origin_masks(p);
		_bit_masks_span = span(_origin, _origin+1);
	}
	assert (!_param.decrease_mask_freedom || !origin_masks()->contour.empty());
	TestCancel();

	calc_data d(*this, request);
	assert(_bit_masks_span.inside(_grey_masks_span));

	d.dir = backward;
	calc(d);
	TestCancel();
	d.dir = forward;
	calc(d);
}

void segment::calc(calc_data & d)
{
	//first of all find grey masks for the frames with bit masks
	int bn = out_of_span(_bit_masks_span, d.dir);
	if (_param.smooth_mask) {
		int gn = out_of_span(_grey_masks_span, d.dir);
		while (bn != gn && !eos(d.request, _grey_masks_span, d.dir)) {
			check_mem();
			twomasks_sptr p(frame_masks(gn)->copy_except_grey());
			icache::cached_image_sptr im_ptr(load_frame(gn, d.dir));
			TestCancel();
			smooth_mask(p, im_ptr->im);
			set_frame_masks(gn, p);
			int n = gn;
			gn = advance(_grey_masks_span, d.dir);
			frame_done(n, d.dir, _grey_masks_span);
			TestCancel();
		}
	}

	if (eos(d.request, _bit_masks_span, d.dir))
		return;

	d.ref.create(d.dir, _bit_masks_span);

	for (;;) {
		check_mem();
		target_frame tag(*this);
		tag.create(bn);
		TestCancel();

		twomasks_sptr p(new twomasks);
		d.ref.track_mask(*p, tag);

		if (_param.dilation_erosion != 0) {
			TestCancel();
			trm::bit_mask b;
			if (_param.dilation_erosion > 0)
				img::dilate4(b, p->b, _param.dilation_erosion);
			else
				img::erode4(b, p->b, -_param.dilation_erosion);
			p->b.swap(b);
		}

		if (_param.wool_removal > 0) {
			TestCancel();
			wool_removal(p->b, _param.wool_removal);
		}

		if (_param.decrease_mask_freedom) {
			TestCancel();

			//usage of very rough p->motion only spoils the contour :(
			//if (p->ref_frame >= 0) {
			//	//apply found by optic flow transformation to reference contour for better initial approximation
			//	bezier_contour etalone;
			//	frame_masks(d.ref.time())->contour.transform(p->motion, etalone);
			//	decrease_mask_freedom(*p, &etalone);
			//}
			//else

				decrease_mask_freedom(*p, &frame_masks(d.ref.time())->contour);
		}

		if (_param.smooth_mask) {
			TestCancel();
			smooth_mask(p, tag.image());
		}

		set_frame_masks(bn, p);

		bn = advance(_bit_masks_span, d.dir);
		advance(_grey_masks_span, d.dir);
		frame_done(tag.time(), d.dir, _param.smooth_mask ? _grey_masks_span : _bit_masks_span);
		if (eos(d.request, _bit_masks_span, d.dir))
			break;

		TestCancel();

		if (std::abs(tag.time() - d.ref.time()) == _param.reference_frame_pitch)
			d.ref.become(tag);
	}
}

void segment::decrease_mask_freedom(twomasks & m, bezier_contour const * etalone)
{
	mswin::timing_mark mark(_timing, "decreasing mask freedom");

	std::vector<float> tractions;
	if (etalone) {
		assert (!etalone->empty());
		m.contour = *etalone;
		tractions.resize(etalone->size());
		unsigned i;
		for (i = 0; i < tractions.size(); ++i)
			tractions[i] = (i % 3 == 0) ? _param.decrease_nodes_traction : _param.decrease_tangents_traction;
	}
	else
		assert (m.contour.empty());

	outline::decrease_mask_freedom(m.b,
		_param.decrease_freedom_degrees,
		_param.decrease_smooth_contour,
		outline::guide_params<float>(
			_param.decrease_iterations,
			0.001f, //stop_mean_discrepancy
			1 / 1.5f, //oversample_factor
			_param.optic_flow ? _param.max_motion : 5 //search_half_window
		),
		m.contour,
		etalone,
		etalone ? &tractions : 0
	);
}

void segment::smooth_mask(twomasks_sptr const & p, byte_image const & im)
{
	mswin::timing_mark mark(_timing, "smoothing mask");

	trm::grey_mask g;

	switch (_param.smooth_mask) {
	/*case smooth_mask_A: {
		using namespace refine::smooth2;
		params pa = default_params;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		pa.fore_depth = _param.smooth_fore_depth;
		pa.back_depth = _param.smooth_back_depth;
		perform(im, p->b, g, pa);
		break;
	}
	case smooth_mask_B: {
		using namespace refine::smooth_b2;
		params pa = default_params;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.fore_depth = _param.smooth_fore_depth;
		pa.back_depth = _param.smooth_back_depth;
		pa.sharpness = _param.smooth_sharpness;
		pa.noisiness = _param.smooth_noisiness;
		perform(im, p->b, g, pa);
		break;
	}*/
	case smooth_mask_C: {
		using namespace refine::smooth_c;
		params pa = default_params;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		//pa.fore_depth = _param.smooth_fore_depth;
		//pa.back_depth = _param.smooth_back_depth;
		pa.fore_depth = pa.back_depth = pa.whs;
		perform(im, p->b, g, pa);
		break;
	}
	/*case smooth_mask_D: {
		using namespace refine::smooth_d;
		params pa = default_params;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		pa.fore_depth = _param.smooth_fore_depth;
		pa.back_depth = _param.smooth_back_depth;
		perform(im, p->b, g, pa);
		break;
	}
	case smooth_mask_E: {
		using namespace refine::smooth_e;
		params pa = default_params;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.fore_depth = _param.smooth_fore_depth;
		pa.back_depth = _param.smooth_back_depth;
		perform(im, p->b, g, pa);
		break;
	}
	case smooth_mask_EA: {
		using namespace refine::smooth_ea;
		params pa = default_params;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		pa.fore_depth = _param.smooth_fore_depth;
		pa.back_depth = _param.smooth_back_depth;
		perform(im, p->b, g, pa);
		break;
	}*/
	case smooth_mask_equal_spheres: {
		using namespace refine::smooth_despill;
		params pa;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		byte_image dup = im;
		//crisp mask to alpha channel of dup
		for (trm::bit_mask::iterator i = p->b.begin(); i != p->b.end(); ++i)
			dup[i].a = *i ? 255 : 0;
		a::perform(dup, pa);
		g = img::projection(dup, &byte_color::a);
		break;
	}
	case smooth_mask_equal_ellipses: {
		using namespace refine::smooth_despill;
		params pa;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		byte_image dup = im;
		//crisp mask to alpha channel of dup
		for (trm::bit_mask::iterator i = p->b.begin(); i != p->b.end(); ++i)
			dup[i].a = *i ? 255 : 0;
		b::perform(dup, pa);
		g = img::projection(dup, &byte_color::a);
		break;
	}
	/*case smooth_mask_different_ellipses: {
		using namespace refine::smooth_despill;
		params pa;
		pa.whs = _param.smooth_window_half_size;
		pa.passes_num = _param.smooth_passes_num;
		pa.max_condition_number = _param.smooth_max_condition_number;
		byte_image dup = im;
		//crisp mask to alpha channel of dup
		for (trm::bit_mask::iterator i = p->b.begin(); i != p->b.end(); ++i)
			dup[i].a = *i ? 255 : 0;
		c::perform(dup, pa);
		g = img::projection(dup, &byte_color::a);
		break;
	}*/
	case smooth_mask_off:
		assert (false);
	}

	if (!_param.smooth_decorative) {
		//binarize grey mask
		for (trm::bit_mask::iterator i = p->b.begin(); i != p->b.end(); ++i)
			*i = (g[i] >= 128);
	}

	p->g.store(g);
}

} //namespace trm
